using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimisers, HDF5
import ModelingToolkit: Interval, infimum, supremum

@parameters x y z t
@variables u(..)
Dt = Differential(t)
Dxx = Differential(x)^2
Dyy = Differential(y)^2
Dzz = Differential(z)^2

# Thermal conductivity constant
k = 1 # replace with actual value

# Gaussian beam source term - time-independent
P = 1 # peak power of laser, replace with actual value
ω = 0.1 # beam waist, replace with actual value
gaussian_beam(x, y, z) = P * exp(-2 * ((x^2 + y^2 + z^2) / ω^2))

# Heat equation (now time-dependent in 3D)
eq = Dt(u(x, y, z, t)) ~ k * (Dxx(u(x, y, z, t)) + Dyy(u(x, y, z, t)) + Dzz(u(x, y, z, t))) - gaussian_beam(x, y, z)

# Boundary and initial conditions
bcs = [u(0, y, z, t) ~ 0.0, u(1, y, z, t) ~ 0,
       u(x, 0, z, t) ~ 0.0, u(x, 1, z, t) ~ 0,
       u(x, y, 0, t) ~ 0.0, u(x, y, 1, t) ~ 0,
       u(x, y, z, 0) ~ 0] # Initial condition at t=0

# Space, depth, and time domains
domains = [x ∈ Interval(0.0, 1.0),
           y ∈ Interval(0.0, 1.0),
           z ∈ Interval(0.0, 1.0), # depth dimension
           t ∈ Interval(0.0, 1.0)] # Adjust time interval as needed

# Discretization
dx = 0.1
dt = 0.1 # time step

# Neural network architecture
dim = 4 # now 4 dimensions including time and z
chain = Lux.Chain(Dense(dim, 16, Lux.σ), Dense(16, 16, Lux.σ), Dense(16, 1))

# Physics-Informed Neural Network
discretization = PhysicsInformedNN(chain, QuadratureTraining())

@named pde_system = PDESystem(eq, bcs, domains, [x, y, z, t], [u(x, y, z, t)])
prob = discretize(pde_system, discretization)

# Callback for loss
callback = function (p, l)
    println("Current loss is: $l")
    return false
end

# Solving the PDE
res = Optimization.solve(prob, ADAM(0.1); callback = callback, maxiters = 4000)
prob = remake(prob, u0 = res.minimizer)
res = Optimization.solve(prob, ADAM(0.01); callback = callback, maxiters = 2000)
phi = discretization.phi

# Prediction and plotting
xs, ys, zs, ts = [infimum(d.domain):(dx / 10):supremum(d.domain) for d in domains]
u_predict = reshape([first(phi([x, y, z, t], res.minimizer)) for x in xs for y in ys for z in zs for t in ts],
                    (length(xs), length(ys), length(zs), length(ts)))

using Plots
# Plot for a specific time slice and z-slice, adjust t_index and z_index as needed
t_index = 1 # Example: first time step
z_index = 1 # Example: first z slice
p3 = plot(xs, ys, u_predict[:,:,z_index,t_index], linetype = :contourf, title = "Predicted at t=$(ts[t_index]), z=$(zs[z_index])")
plot(p3)


# Assuming u_predict is already computed as in your code

#anim = @animate for t_index in 1:length(ts)
#    p = plot(layout = (1,1), size = (600, 600))
#    for z_index in 1:length(zs)
#        # Creating a semi-transparent layer for each z-slice
#        contourf!(p, xs, ys, u_predict[:,:,z_index,t_index], alpha = 0.5, title = "Time: $(ts[t_index])")
#    end
#end

# Save the animation
#gif(anim, "heat_distribution_evolution.gif", fps = 15)

# Save the array to an HDF5 file
h5write("data.h5", "/u_predict", u_predict)
